import torch
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import get_image_colorizer
import gradio as gr

# 检查并设置设备
if torch.cuda.is_available():
    print("CUDA is available. Using GPU.")
    device.set(device=DeviceId.GPU0)  # 设置为第一个GPU
else:
    print("CUDA is not available. Using CPU.")
    device.set(device=DeviceId.CPU)

# 加载预训练的 DeOldify 模型
colorizer = get_image_colorizer(artistic=True)

def colorize_image(input_image):
    # 使用 DeOldify 模型将输入图像进行彩色化
    colorized_image = colorizer.get_transformed_image(input_image, render_factor=35)
    return colorized_image

# 创建 Gradio 界面
interface = gr.Interface(
    fn=colorize_image,
    inputs=gr.inputs.Image(type="pil", label="Upload Black and White Image"),
    outputs=gr.outputs.Image(type="pil", label="Colorized Image"),
    title="Black and White to Color Image",
    description="Upload a black and white image to convert it into a color image using DeOldify."
)

# 输出当前设备信息
print("Current device:", device)

# 启动 Gradio 应用
interface.launch()
